"""
Differential operators for cylindrically symmetric grids.

This module implements differential operators for cylindrical coordinate systems
with cylindrical symmetry (r, z dependence only, no angular θ dependence).

Mathematical forms:
- Laplacian: ∇²u = ∂²u/∂r² + (1/r)∂u/∂r + ∂²u/∂z²
- Gradient: ∇u = (∂u/∂r, ∂u/∂z)
- Divergence: ∇·v = (1/r)∂(rvᵣ)/∂r + ∂vᵤ/∂z
- Vector Laplacian: ∇²v = (∇²vᵣ - vᵣ/r², ∇²vᵤ)

Key features:
- Proper handling of r=0 singularity on the axis
- Conservative and non-conservative formulations
- Support for mixed boundary conditions in r and z directions
- Optimized for axisymmetric problems
"""

using LinearAlgebra
using SparseArrays

include("common.jl")

# Export cylindrical operator functions
export make_laplace_cylindrical, make_gradient_cylindrical, make_divergence_cylindrical
export make_vector_laplace_cylindrical, make_poisson_solver_cylindrical
export make_conservative_laplace_cylindrical

"""
    make_laplace_cylindrical(grid::CylindricalSymGrid; conservative=false) -> OperatorType

Create Laplacian operator for cylindrically symmetric grid.

The cylindrical Laplacian is:
∇²u = ∂²u/∂r² + (1/r)∂u/∂r + ∂²u/∂z²

At r=0, using L'Hôpital's rule: (1/r)∂u/∂r → ∂²u/∂r²
So the effective r-component at r=0 becomes: 2∂²u/∂r²

Args:
    grid: Cylindrically symmetric grid
    conservative: Whether to use conservative discretization

Returns:
    Function that applies cylindrical Laplacian
"""
function make_laplace_cylindrical(grid::CylindricalSymGrid; conservative=false)
    if conservative
        return make_conservative_laplace_cylindrical(grid)
    else
        return make_standard_laplace_cylindrical(grid)
    end
end

"""
    make_standard_laplace_cylindrical(grid::CylindricalSymGrid) -> OperatorType

Create standard cylindrical Laplacian operator.
"""
function make_standard_laplace_cylindrical(grid::CylindricalSymGrid)
    nr, nz = grid.shape
    dr, dz = spacing(grid)
    rs = coordinate_arrays(grid)[1]  # radial coordinates
    zs = coordinate_arrays(grid)[2]  # axial coordinates  
    r_min = extent(grid)[1][1]
    
    dr2_inv = 1.0 / (dr * dr)
    dz2_inv = 1.0 / (dz * dz)
    dr_inv = 1.0 / (2.0 * dr)
    dz_inv = 1.0 / (2.0 * dz)
    
    function laplace_cylindrical(field_data, out)
        for ir in 1:nr, iz in 1:nz
            r = rs[ir]
            
            # z-direction second derivative (same for all r)
            d2u_dz2 = if iz == 1 || iz == nz
                0.0  # Boundary points handled separately
            else
                (field_data[ir, iz+1] - 2*field_data[ir, iz] + field_data[ir, iz-1]) * dz2_inv
            end
            
            # r-direction derivatives
            if ir == 1 && abs(r_min) < 1e-12  # r=0 axis
                # At r=0: ∇²u = 4∂²u/∂r² + ∂²u/∂z² using L'Hôpital's rule
                d2u_dr2 = if iz == 1 || iz == nz || nr < 2
                    0.0  # Not enough points or boundary
                else
                    4.0 * (field_data[ir+1, iz] - field_data[ir, iz]) * dr2_inv
                end
                
                out[ir, iz] = d2u_dr2 + d2u_dz2
                
            elseif ir == nr  # Outer boundary
                out[ir, iz] = 0.0  # Will be set by boundary conditions
                
            else  # Interior points in r
                d2u_dr2 = (field_data[ir+1, iz] - 2*field_data[ir, iz] + field_data[ir-1, iz]) * dr2_inv
                du_dr = (field_data[ir+1, iz] - field_data[ir-1, iz]) * dr_inv
                
                out[ir, iz] = d2u_dr2 + (1.0/r) * du_dr + d2u_dz2
            end
        end
        
        return out
    end
end

"""
    make_conservative_laplace_cylindrical(grid::CylindricalSymGrid) -> OperatorType

Create conservative cylindrical Laplacian operator.

The conservative form discretizes:
∇²u = (1/r)∂/∂r(r∂u/∂r) + ∂²u/∂z²
"""
function make_conservative_laplace_cylindrical(grid::CylindricalSymGrid)
    nr, nz = grid.shape
    dr, dz = spacing(grid)
    rs = coordinate_arrays(grid)[1]
    r_min, r_max = extent(grid)[1]
    
    # Create cell boundaries for conservative discretization
    rl = rs .- dr/2  # inner radii
    rh = rs .+ dr/2  # outer radii
    rl[1] = r_min
    rh[end] = r_max
    
    dz2_inv = 1.0 / (dz * dz)
    
    function conservative_laplace_cylindrical(field_data, out)
        for ir in 1:nr, iz in 1:nz
            r = rs[ir]
            
            # z-direction second derivative (conservative form is same as standard)
            d2u_dz2 = if iz == 1 || iz == nz
                0.0  # Boundary
            else
                (field_data[ir, iz+1] - 2*field_data[ir, iz] + field_data[ir, iz-1]) * dz2_inv
            end
            
            # r-direction conservative derivative: (1/r)∂/∂r(r∂u/∂r)
            if ir == 1  # r=0 axis
                # Flux from axis is zero due to symmetry
                if ir < nr
                    flux_outer = rh[ir] * (field_data[ir+1, iz] - field_data[ir, iz]) / dr
                    radial_term = flux_outer / (r * dr)
                else
                    radial_term = 0.0
                end
                
            elseif ir == nr  # Outer boundary  
                flux_inner = rl[ir] * (field_data[ir, iz] - field_data[ir-1, iz]) / dr
                radial_term = -flux_inner / (r * dr)
                
            else  # Interior points
                flux_outer = rh[ir] * (field_data[ir+1, iz] - field_data[ir, iz]) / dr
                flux_inner = rl[ir] * (field_data[ir, iz] - field_data[ir-1, iz]) / dr
                radial_term = (flux_outer - flux_inner) / (r * dr)
            end
            
            out[ir, iz] = radial_term + d2u_dz2
        end
        
        return out
    end
end

"""
    make_gradient_cylindrical(grid::CylindricalSymGrid, dim::Int; method="central") -> OperatorType

Create gradient operator for cylindrical grid.

Args:
    grid: Cylindrical grid
    dim: Dimension (1 for r, 2 for z)  
    method: Finite difference method

Returns:
    Function that computes gradient component
"""
function make_gradient_cylindrical(grid::CylindricalSymGrid, dim::Int; method="central")
    nr, nz = grid.shape
    dr, dz = spacing(grid)
    
    if dim == 1  # ∂/∂r
        if method == "central"
            dr_inv = 1.0 / (2.0 * dr)
            
            function gradient_r_central(field_data, out)
                for ir in 1:nr, iz in 1:nz
                    if ir == 1
                        # Forward difference at r=0
                        out[ir, iz] = (field_data[ir+1, iz] - field_data[ir, iz]) / dr
                    elseif ir == nr
                        # Backward difference at boundary
                        out[ir, iz] = (field_data[ir, iz] - field_data[ir-1, iz]) / dr
                    else
                        # Central difference
                        out[ir, iz] = (field_data[ir+1, iz] - field_data[ir-1, iz]) * dr_inv
                    end
                end
                return out
            end
        else
            error("Only central differences implemented for radial gradient")
        end
        
    elseif dim == 2  # ∂/∂z
        if method == "central"
            dz_inv = 1.0 / (2.0 * dz)
            
            function gradient_z_central(field_data, out)
                for ir in 1:nr, iz in 1:nz
                    if iz == 1
                        # Forward difference at z boundary
                        out[ir, iz] = (field_data[ir, iz+1] - field_data[ir, iz]) / dz
                    elseif iz == nz
                        # Backward difference at z boundary  
                        out[ir, iz] = (field_data[ir, iz] - field_data[ir, iz-1]) / dz
                    else
                        # Central difference
                        out[ir, iz] = (field_data[ir, iz+1] - field_data[ir, iz-1]) * dz_inv
                    end
                end
                return out
            end
        else
            error("Only central differences implemented for axial gradient")
        end
    else
        error("Invalid dimension for cylindrical gradient: $dim")
    end
end

"""
    make_divergence_cylindrical(grid::CylindricalSymGrid) -> OperatorType

Create divergence operator for cylindrical coordinates.

∇·v = (1/r)∂(rvᵣ)/∂r + ∂vᵤ/∂z
"""
function make_divergence_cylindrical(grid::CylindricalSymGrid)
    nr, nz = grid.shape
    dr, dz = spacing(grid)
    rs = coordinate_arrays(grid)[1]
    
    dr_inv = 1.0 / (2.0 * dr)
    dz_inv = 1.0 / (2.0 * dz)
    
    function divergence_cylindrical(vector_field, out)
        # vector_field should be (nr, nz, 2) with components [vr, vz]
        
        for ir in 1:nr, iz in 1:nz
            r = rs[ir]
            
            # ∂vᵤ/∂z term
            dvz_dz = if iz == 1 || iz == nz
                0.0  # Boundary handling
            else
                (vector_field[ir, iz+1, 2] - vector_field[ir, iz-1, 2]) * dz_inv
            end
            
            # (1/r)∂(rvᵣ)/∂r term
            if ir == 1  # r=0 axis
                # Using L'Hôpital's rule: (1/r)∂(rvᵣ)/∂r → 2∂vᵣ/∂r at r=0
                if ir < nr
                    dvr_dr_times_2 = 2.0 * (vector_field[ir+1, iz, 1] - vector_field[ir, iz, 1]) / dr
                    out[ir, iz] = dvr_dr_times_2 + dvz_dz
                else
                    out[ir, iz] = dvz_dz
                end
                
            elseif ir == nr  # Outer boundary
                out[ir, iz] = 0.0  # Boundary handling
                
            else  # Interior points
                # ∂(rvᵣ)/∂r using central differences
                rvr_forward = (rs[ir] + dr/2) * (vector_field[ir+1, iz, 1] + vector_field[ir, iz, 1])/2
                rvr_backward = (rs[ir] - dr/2) * (vector_field[ir, iz, 1] + vector_field[ir-1, iz, 1])/2
                d_rvr_dr = (rvr_forward - rvr_backward) / dr
                
                out[ir, iz] = d_rvr_dr/r + dvz_dz
            end
        end
        
        return out
    end
end

"""
    make_vector_laplace_cylindrical(grid::CylindricalSymGrid) -> OperatorType

Create vector Laplacian for cylindrical coordinates.

∇²v = (∇²vᵣ - vᵣ/r², ∇²vᵤ)êᵣ + ∇²vᵤêᵤ

Note: The angular component is zero in cylindrical symmetry.
"""
function make_vector_laplace_cylindrical(grid::CylindricalSymGrid)
    scalar_laplace = make_laplace_cylindrical(grid)
    nr, nz = grid.shape
    rs = coordinate_arrays(grid)[1]
    
    function vector_laplace_cylindrical(vector_field, out)
        # Process radial component: ∇²vᵣ - vᵣ/r²
        vr = view(vector_field, :, :, 1)
        temp_r = similar(vr)
        scalar_laplace(vr, temp_r)
        
        for ir in 1:nr, iz in 1:nz
            r = rs[ir]
            if r > 1e-12
                out[ir, iz, 1] = temp_r[ir, iz] - vr[ir, iz] / (r * r)
            else
                # At r=0, the correction term vanishes
                out[ir, iz, 1] = temp_r[ir, iz]
            end
        end
        
        # Process axial component: ∇²vᵤ (no correction term)
        vz = view(vector_field, :, :, 2)
        temp_z = view(out, :, :, 2)
        scalar_laplace(vz, temp_z)
        
        return out
    end
end

"""
    make_poisson_solver_cylindrical(grid::CylindricalSymGrid, bcs) -> Function

Create Poisson solver for cylindrical grid: ∇²u = f
"""
function make_poisson_solver_cylindrical(grid::CylindricalSymGrid, bcs)
    return make_general_poisson_solver(grid, bcs)
end

"""
    make_laplace_matrix_cylindrical(grid::CylindricalSymGrid, bcs) -> Tuple{SparseMatrixCSC, Vector}

Create sparse Laplacian matrix for cylindrical grid.
"""
function make_laplace_matrix_cylindrical(grid::CylindricalSymGrid, bcs)
    nr, nz = grid.shape
    dr, dz = spacing(grid)
    rs = coordinate_arrays(grid)[1]
    r_min = extent(grid)[1][1]
    
    total_points = nr * nz
    I_vals = Int[]
    J_vals = Int[]
    V_vals = Float64[]
    rhs_vector = zeros(total_points)
    
    # Linear index helper function
    idx(ir, iz) = (ir - 1) * nz + iz
    
    dr2_inv = 1.0 / (dr * dr)
    dz2_inv = 1.0 / (dz * dz)
    dr_inv = 1.0 / (2.0 * dr)
    
    for ir in 1:nr, iz in 1:nz
        i = idx(ir, iz)
        r = rs[ir]
        
        # Handle z-direction derivatives
        if iz == 1  # z boundary (bottom)
            # Apply boundary condition
            bc_z = bcs[2]  # z-direction boundary condition
            if isa(bc_z, DirichletBC)
                push!(I_vals, i); push!(J_vals, i); push!(V_vals, 1.0)
                rhs_vector[i] += bc_z.value
            end
        elseif iz == nz  # z boundary (top)
            # Apply boundary condition  
            bc_z = bcs[2]
            if isa(bc_z, DirichletBC)
                push!(I_vals, i); push!(J_vals, i); push!(V_vals, 1.0)
                rhs_vector[i] += bc_z.value
            end
        else  # Interior z points
            # Second derivative in z: ∂²u/∂z²
            push!(I_vals, i); push!(J_vals, idx(ir, iz-1)); push!(V_vals, dz2_inv)
            push!(I_vals, i); push!(J_vals, i); push!(V_vals, -2.0 * dz2_inv)
            push!(I_vals, i); push!(J_vals, idx(ir, iz+1)); push!(V_vals, dz2_inv)
        end
        
        # Handle r-direction derivatives
        if ir == 1 && abs(r_min) < 1e-12  # r=0 axis
            # Special stencil: 4∂²u/∂r²
            if nr > 1 && iz > 1 && iz < nz  # Interior in z
                push!(I_vals, i); push!(J_vals, i); push!(V_vals, -4.0 * dr2_inv)
                push!(I_vals, i); push!(J_vals, idx(ir+1, iz)); push!(V_vals, 4.0 * dr2_inv)
            end
        elseif ir == nr  # Outer r boundary
            bc_r = bcs[1]  # r-direction boundary condition
            if isa(bc_r, DirichletBC)
                push!(I_vals, i); push!(J_vals, i); push!(V_vals, 1.0)
                rhs_vector[i] += bc_r.value
            elseif isa(bc_r, NeumannBC)
                push!(I_vals, i); push!(J_vals, i); push!(V_vals, 1.0/dr)
                push!(I_vals, i); push!(J_vals, idx(ir-1, iz)); push!(V_vals, -1.0/dr)
                rhs_vector[i] += bc_r.value
            end
        else  # Interior r points
            if iz > 1 && iz < nz  # Also interior in z
                # Standard cylindrical Laplacian stencil
                # ∂²u/∂r² term
                push!(I_vals, i); push!(J_vals, idx(ir-1, iz)); push!(V_vals, dr2_inv)
                push!(I_vals, i); push!(J_vals, i); push!(V_vals, -2.0 * dr2_inv)
                push!(I_vals, i); push!(J_vals, idx(ir+1, iz)); push!(V_vals, dr2_inv)
                
                # (1/r)∂u/∂r term
                push!(I_vals, i); push!(J_vals, idx(ir-1, iz)); push!(V_vals, -dr_inv/r)
                push!(I_vals, i); push!(J_vals, idx(ir+1, iz)); push!(V_vals, dr_inv/r)
            end
        end
    end
    
    matrix = sparse(I_vals, J_vals, V_vals, total_points, total_points)
    return matrix, rhs_vector
end

# Make matrix function available to general solver
make_laplace_matrix(grid::CylindricalSymGrid, bcs) = make_laplace_matrix_cylindrical(grid, bcs)

"""
    cylindrical_volume_element(r, dr, dz) -> Float64

Calculate volume element for cylindrical coordinates: 2πr*dr*dz
"""
function cylindrical_volume_element(r, dr, dz)
    return 2π * r * dr * dz
end

"""
    apply_cylindrical_bc!(operator_matrix, rhs_vector, bcs, grid)

Apply boundary conditions for cylindrical grids.
"""
function apply_cylindrical_bc!(operator_matrix, rhs_vector, bcs, grid)
    nr, nz = grid.shape
    
    # Apply r-direction boundary conditions
    bc_r = bcs[1]
    for iz in 1:nz
        i = (nr - 1) * nz + iz  # Last r index
        
        if isa(bc_r, DirichletBC)
            operator_matrix[i, :] .= 0.0
            operator_matrix[i, i] = 1.0
            rhs_vector[i] = bc_r.value
        elseif isa(bc_r, NeumannBC)
            # Modify for gradient BC
            dr = spacing(grid)[1]
            operator_matrix[i, :] .= 0.0
            operator_matrix[i, i] = 1.0/dr
            operator_matrix[i, i-nz] = -1.0/dr
            rhs_vector[i] = bc_r.value
        end
    end
    
    # Apply z-direction boundary conditions
    bc_z = bcs[2]
    for ir in 1:nr
        # Bottom boundary (iz = 1)
        i = (ir - 1) * nz + 1
        if isa(bc_z, DirichletBC)
            operator_matrix[i, :] .= 0.0
            operator_matrix[i, i] = 1.0
            rhs_vector[i] = bc_z.value
        end
        
        # Top boundary (iz = nz)  
        i = (ir - 1) * nz + nz
        if isa(bc_z, DirichletBC)
            operator_matrix[i, :] .= 0.0
            operator_matrix[i, i] = 1.0
            rhs_vector[i] = bc_z.value
        end
    end
    
    return operator_matrix, rhs_vector
end